#include <iostream>
#include <stdlib.h>
#include <random>
#include <cmath>

#include "fmt/core.h"
#include "fmt/ranges.h"
#include "fmt/color.h"
#include "gtest/gtest.h"

#ifdef USE_OPEN_CV
#include "opencv2/opencv.hpp"
#include "opencv2/core/utils/logger.hpp"
#endif

#include "KpExtGP.h"

using namespace std;
using  chronoms = std::chrono::milliseconds;

void msleep(int ms){
    std::this_thread::sleep_for(
                chronoms(ms)
                );
}

static bool CtxInitOk = false;

TEST(KpExt, Initial){
    if(kpextInit() == 0){
        CtxInitOk = true;
    }
    ASSERT_EQ(true, CtxInitOk);
}


//  图像处理阶段1 测试
TEST(KpExt, ImgProcessStage1){
    ASSERT_TRUE(CtxInitOk);
    int cols_in = 10, rows_in = 10;
    int cols_out, rows_out;
    int avgcnt = 1;

    {  //  基本流程测试
        kpextCfgStage1Input(cols_in, rows_in, 0);
        kpextCfgStage1Win(0);
        kpextCfgStage1Amp(16, 0, true);
        kpextCfgStage1Avg(avgcnt);
        kpextCfgStage1ROI(0, 0, 0, 0);
        kpextCfgStage1Flip(false, false);
        kpextApplyCfgStage1();

        kpextGetStage1OutputSize(&cols_out, &rows_out);

        kpextCfgStage1RingBufferCnt(100);

        kpextStartStage1();
        size_t size = static_cast<size_t>(cols_in * rows_in);
        auto data = std::shared_ptr<uint16_t>(new uint16_t[size]);
        for (int i = 0; i < 100; ++i) {
            data.get()[i] = static_cast<uint16_t>(i);
        }

        int avgidx = 0;
        int timeout = 5000;
        for (int i = 0; i < 100; ++i) {
            data.get()[0] = static_cast<uint16_t>(i);
            kpextEnqueStage1(data.get(), nullptr);
            if(avgidx == avgcnt - 1){
                auto ptr = kpextWaitStage1Output(timeout, nullptr);
                if(i < 10 && ptr){
                    timeout = 100;
                    std::vector<uint16_t> disp(ptr, ptr + cols_out * rows_out);
                    fmt::print("Output : {}\n", disp);
                }
            }
            avgidx++;
            avgidx = avgidx % avgcnt;
        }

        kpextStopStage1();
    }

    std::cout<< "\n ********************************************* \n";

    {  //  基本用时测试
        cols_in = 1024;
        rows_in = 500;
        avgcnt = 1;

        kpextCfgStage1Input(cols_in, rows_in, 0);
        kpextCfgStage1DCResample(0.01f);
        kpextCfgStage1Win(1);
        kpextCfgStage1Amp(1, 0, false);
        kpextCfgStage1Avg(avgcnt);
        kpextCfgStage1ROI(0, 0, 0, 0);
        kpextCfgStage1Flip(true, false);

        kpextApplyCfgStage1();
        kpextGetStage1OutputSize(&cols_out, &rows_out);

        kpextCfgStage1RingBufferCnt(0);

        kpextStartStage1();

        size_t size = static_cast<size_t>(cols_in * rows_in);
        std::shared_ptr<uint16_t> datas[100];
        for (int i = 0; i < 100; ++i) {
            datas[i] = std::shared_ptr<uint16_t>(new uint16_t[size]);
        }

        double pt = 0;
        int timeout = 5000;
        int outcnt = 0;

        for (int i = 0; i < avgcnt; ++i){
            kpextEnqueStage1(datas[i].get(), datas[i].get());
        }
        kpextWaitStage1Output(timeout, nullptr);

        auto begin = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < 100; ++i){
            kpextEnqueStage1(datas[i].get(), datas[i].get());
        }
        while (outcnt < (100 / avgcnt)) {
            auto ptr = kpextWaitStage1Output(1, nullptr);
            if(ptr){
                outcnt ++;
            }
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
        pt = static_cast<double>(elapsed.count());
        pt = pt * 1e-6;

        kpextStopStage1();

        fmt::print("\n *********************************************\n");
        fmt::print("ImgProcessStage1 benchmark --> \n"
                   "Frame size : ({}x{}),\n"
                   "{} cycle,\n"
                   "total time : {}ms,\n"
                   "each FFT time : {}ms",
                   rows_in, cols_in,
                   100, pt,
                   pt / 100
                   );
        fmt::print("\n *********************************************\n");
    }
}

//  图像处理阶段2 测试
TEST(KpExt, ImgProcessStage2){
    ASSERT_TRUE(CtxInitOk);

    std::cout<< "Press Enter to go on...\n";
    std::cin.get();

    int cols_in = 10, rows_in = 10;
    int cols_out, rows_out;
    {  //  基本流程测试
        kpextSetStage2CMap(0);
        kpextSetStage2HFlip(false);
        kpextSetStage2VFlip(false);
        kpextSetStage2Threshold(0);
        kpextSetStage2Mul(1);
        kpextSetStage2Power(1);

        kpextStartStage2();

        size_t size = static_cast<size_t>(cols_in * rows_in);
        auto data = std::shared_ptr<uint16_t>(new uint16_t[size]);
        std::shared_ptr<uint8_t> outputs[4];
        for (int i = 0; i < 4; ++i) {
            outputs[i] = std::shared_ptr<uint8_t>(new uint8_t[4 * size]);
        }
        for (int i = 0; i < 100; ++i) {
            data.get()[i] = static_cast<uint16_t>(i * 256);
        }

        int timeout = 5000;
        float a = 1;
        for (int i = 0; i < 100; ++i) {
            std::memset(outputs[i % 4].get(), 0, size * 4);

            data.get()[0] = static_cast<uint16_t>(i * 256);
            kpextEnqueStage2(cols_in, rows_in,
                             data.get(),
                             outputs[i % 4].get()
                             );

            auto ptr = kpextWaitStage2Output(&cols_out, &rows_out, timeout);
            if(i < 10 && ptr){
                timeout = 100;
                std::vector<uint8_t> disp(ptr, ptr + cols_out * rows_out * 4);
                fmt::print("Output : {}\n", disp);
//                break;
            }
            a = a + 0.1f;
            kpextSetStage2Mul(a);
        }

        kpextStopStage2();
    }

    std::cout<< "\n ********************************************* \n";

    {  //  基本用时测试
        cols_in = 500;
        rows_in = 512;

        kpextSetStage2CMap(0);
        kpextSetStage2HFlip(true);
        kpextSetStage2VFlip(false);
        kpextSetStage2Threshold(10000);
        kpextSetStage2Mul(1);
        kpextSetStage2Power(1);

        kpextStartStage2();

        size_t size = static_cast<size_t>(cols_in * rows_in);
        auto data = std::shared_ptr<uint16_t>(new uint16_t[size]);
        std::shared_ptr<uint8_t> outputs[100];
        for (int i = 0; i < 100; ++i) {
            outputs[i] = std::shared_ptr<uint8_t>(new uint8_t[size * 4]);
        }

        double pt = 0;
        int timeout = 5000;
        int outcnt = 0;

        kpextEnqueStage2(cols_in, rows_in,
                         data.get(),
                         outputs[0].get()
                         );
        kpextWaitStage2Output(&cols_out, &rows_out, timeout);
        auto begin = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < 100; ++i){
            kpextEnqueStage2(cols_in, rows_in,
                             data.get(),
                             outputs[i].get()
                             );
        }
        while (outcnt < 100) {
            auto ptr = kpextWaitStage2Output(&cols_out, &rows_out, 1);
            if(ptr){
                outcnt ++;
            }
        }

        auto end = std::chrono::high_resolution_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
        pt = static_cast<double>(elapsed.count());
        pt = pt * 1e-6;

        kpextStopStage2();

        fmt::print("\n *********************************************\n");
        fmt::print("ImgProcessStage2 benchmark --> \n"
                   "Frame size : ({}x{}),\n"
                   "{} cycle,\n"
                   "total time : {}ms,\n"
                   "each ColorMap time : {}ms",
                   rows_in, cols_in,
                   100, pt,
                   pt / 100
                   );
        fmt::print("\n *********************************************\n");
    }
}

#ifdef USE_OPEN_CV
TEST(KpExt, ImgProcessStages_CV){
    ASSERT_TRUE(CtxInitOk);

    //  图像处理阶段1  可视化测试
    std::cout<< "\n ********************************************* \n";
    std::cout<< "Press Enter to go on...\n";
    std::cin.get();

    int cols_in = 1024;
    int rows_in = 500;
    int avgcnt = 1;
    int cols_out, rows_out;

    kpextCfgStage1Input(cols_in, rows_in, 0);
    kpextCfgStage1DCResample(0.01f);
    kpextCfgStage1Win(1);
    kpextCfgStage1Amp(7908, 0, false); //static_cast<float>(cols_in)
    kpextCfgStage1Avg(avgcnt);
    kpextCfgStage1ROI(0, 0, 0, 0);
    kpextCfgStage1Flip(false, true);

    kpextApplyCfgStage1();
    kpextGetStage1OutputSize(&cols_out, &rows_out);

    kpextCfgStage1RingBufferCnt(0);
    kpextStartStage1();

    size_t size = static_cast<size_t>(cols_in * rows_in);
    auto data = std::shared_ptr<uint16_t>(new uint16_t[size]);
    std::memset(data.get(), 0, size * 2);
    for (int r = 0; r < rows_in; ++r) {
        int idx = r * cols_in + cols_in / 2;
        data.get()[idx] = 65535;
        data.get()[idx + 1] = 65535;
        data.get()[idx + 2] = 65535;
    }

    std::shared_ptr<uint16_t> outputs[4];
    for (int i = 0; i < 4; ++i) {
        outputs[i] = std::shared_ptr<uint16_t>(new uint16_t[size]);
    }

    cv::Mat mat;
    int timeout = 5000;
    for (int i = 0; i < 16; ++i){
        bool exit = false;
        kpextEnqueStage1(data.get(), outputs[i % 4].get());
        if(i % 4 == 3){
            for (int j = 0; j < 4; ++j) {
                auto ptr = kpextWaitStage1Output(timeout, nullptr);
                if(ptr){
                    mat = cv::Mat(rows_out, cols_out, CV_16UC1, ptr);
                    cv::imshow("ImgProcessStage1_MeanConv", mat);
                    exit = cv::waitKey(1) == 27;
                    timeout = 100;
                }else{
                    std::cout<< "\nkpextWaitStage1Output time out!\n";
                }
            }
        }
        if(exit) break;
    }

    kpextStopStage1();
    kpextStartStage1();

    {
        const float cmlut[]={
            0,   0,0,0,255,
            103, 103.0f,0,0,255,
            104, 104.0f,0,0,255,
            142, 142.0f,64.0f,0,255,
            143, 143.0f,65.0f,2.0f,255,
            239, 239.0f,227.0f,218.0f,255,
            240, 240.0f,229.0f,221.0f,255,
            255, 255,255,255,255,
        };

        kpextCfgStage2CMapLUT(1, cmlut, 8);
    }
    kpextStartStage2();

    std::cout<< "\n ********************************************* \n";
    std::random_device r;
    std::default_random_engine e1(r());
    std::uniform_int_distribution<uint16_t> uniform_dist(0, 65535);

    std::shared_ptr<uint8_t> coloroutputs[4];
    for (int i = 0; i < 4; ++i) {
        coloroutputs[i] = std::shared_ptr<uint8_t>(new uint8_t[size * 4]);
    }

    cv::Mat cmat;
    kpextEnqueStage1(data.get(), outputs[0].get());
    kpextWaitStage1Output(5000, nullptr);
    for (int i = 0; i < 100; ++i){
        bool exit = false;
        kpextEnqueStage1(data.get(), outputs[i % 4].get());
        auto ptr = kpextWaitStage1Output(timeout, nullptr);
        if(ptr){
            mat = cv::Mat(rows_out, cols_out, CV_16UC1, ptr);
            cv::imshow("ImgProcessStage1_RandConv", mat);
            exit = cv::waitKey(33) == 27;
//                std::memset(data.get(), 0, size * 2);
            if(exit) break;

            for (int r = 0; r < rows_in; ++r) {
                int idx = r * cols_in + cols_in / 2;
                uint16_t ttt[3]={uniform_dist(e1), uniform_dist(e1), uniform_dist(e1)};
                data.get()[idx] = ttt[0];
                data.get()[idx + 1] = ttt[1];
                data.get()[idx + 2] = ttt[2];
            }

            if(uniform_dist(e1) % 8 == 0){
                kpextSetStage2CMap(uniform_dist(e1) % 2 == 0);
            }
            kpextSetStage2HFlip(uniform_dist(e1) % 2 == 0);
            kpextSetStage2VFlip(uniform_dist(e1) % 4 == 0);
            kpextSetStage2Threshold(uniform_dist(e1) / 100);
            kpextSetStage2Mul(uniform_dist(e1) / 31845.0f); //uniform_dist(e1) / 50000.0f
            kpextSetStage2Power(1.0f + uniform_dist(e1) / 12e5f);

            int c = cols_out, r = rows_out;
            if(uniform_dist(e1) % 20 == 0){
                c = c / 2;
            }
            if(uniform_dist(e1) % 16 == 0){
                r = r / 2;
            }

            kpextEnqueStage2(c, r,
                             ptr,
                             coloroutputs[i % 4].get()
                    );

            {
                int c, r;
                auto ptr1 = kpextWaitStage2Output(&c, &r, timeout);
                if(ptr1){
                    cmat = cv::Mat(r, c, CV_8UC4, ptr1);
                    cv::cvtColor(cmat, cmat, cv::COLOR_RGBA2BGRA);
                    cv::imshow("ImgProcessStage2_ColorMap", cmat);
                    exit = cv::waitKey(33) == 27;
                    if(exit) break;
                }else{
                    std::cout<< "\nkpextWaitStage2Output time out!\n";
                }
            }
        }else{
            std::cout<< "\nkpextWaitStage1Output time out!\n";
        }
    }

    kpextStopStage1();
    kpextStopStage2();
    cv::destroyAllWindows();
}
#endif

TEST(KpExt, Release){
    ASSERT_EQ(0, kpextRelease());
}

GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
#   ifdef USE_OPEN_CV
    cv::utils::logging::setLogLevel(
                cv::utils::logging::LOG_LEVEL_WARNING);  //  警告级以上信息才显示
#   endif

    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");
    //    kpextSetVKValidationLayerEn(true);
    //    kpextSetVKComputeShaderDefWGS(256,2,3);

    //    kpextSetStagesThreadsCount(4, 1);

    //  _putenv("GTEST_BREAK_ON_FAILURE=1");
    testing::InitGoogleTest(&argc, argv);

    //    testing::Environment* const foo_env =
    //        testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    //    UNUSE_VAR(foo_env);

    //  ::testing::FLAGS_gtest_break_on_failure = true;  //  出错就退出整个测试程序
    return RUN_ALL_TESTS();
}
